package network

import (
	"gonum.org/v1/gonum/mat"
)

// FlattenLayer 实现扁平化层，将3D张量转换为向量
type FlattenLayer struct {
	InputDepth  int
	InputHeight int
	InputWidth  int
	OutputSize  int
	InputShape  []int
}

// NewFlattenLayer 创建一个新的扁平化层
func NewFlattenLayer(inputDepth, inputHeight, inputWidth int) *FlattenLayer {
	outputSize := inputDepth * inputHeight * inputWidth

	return &FlattenLayer{
		InputDepth:  inputDepth,
		InputHeight: inputHeight,
		InputWidth:  inputWidth,
		OutputSize:  outputSize,
		InputShape:  []int{inputDepth, inputHeight, inputWidth},
	}
}

// Forward 实现扁平化层的前向传播
func (f *FlattenLayer) Forward(input interface{}) interface{} {
	inputTensor, ok := input.(*Tensor3D)
	if !ok {
		panic("FlattenLayer.Forward: input is not a *Tensor3D")
	}

	// 创建输出向量
	output := mat.NewVecDense(f.OutputSize, nil)

	// 将3D张量扁平化为1D向量
	idx := 0
	for d := 0; d < inputTensor.Depth; d++ {
		for h := 0; h < inputTensor.Height; h++ {
			for w := 0; w < inputTensor.Width; w++ {
				output.SetVec(idx, inputTensor.Data[d][h][w])
				idx++
			}
		}
	}

	return output
}

// Backward 实现扁平化层的反向传播
func (f *FlattenLayer) Backward(outputGradient interface{}, learningRate float64) interface{} {
	gradVec, ok := outputGradient.(*mat.VecDense)
	if !ok {
		panic("FlattenLayer.Backward: outputGradient is not a *mat.VecDense")
	}

	// 创建输入梯度张量
	inputGrad := NewTensor3D(f.InputDepth, f.InputHeight, f.InputWidth)

	// 将1D向量转换回3D张量
	idx := 0
	for d := 0; d < f.InputDepth; d++ {
		for h := 0; h < f.InputHeight; h++ {
			for w := 0; w < f.InputWidth; w++ {
				inputGrad.Data[d][h][w] = gradVec.AtVec(idx)
				idx++
			}
		}
	}

	return inputGrad
}

// GetType 返回层的类型
func (f *FlattenLayer) GetType() string {
	return "Flatten"
}

// GetOutputShape 返回输出形状
func (f *FlattenLayer) GetOutputShape() []int {
	return []int{f.OutputSize}
}
